package com.eonmind.ai.controller;

import com.eonmind.ai.service.RagService;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import reactor.core.publisher.Flux;

/**
 * Rag AI问答
 */
//@RestController
//@RequestMapping("/ai")
public class CloudRagController {

    private final RagService cloudRagService;

    public CloudRagController(RagService cloudRagService, QueryTransformer queryTransformer) {
        this.cloudRagService = cloudRagService;
        this.queryTransformer = queryTransformer;
    }

    /**
     * 导入文档
     */
    @GetMapping("/bailian/knowledge/importDocument")
    public void importDocument() {
        cloudRagService.importDocuments();
    }

    private final QueryTransformer queryTransformer;

    /**
     * 知识库问答
     *
     * @param message
     * @param response
     * @return
     */
    @GetMapping("/bailian/knowledge/generate")
    public Flux<String> generate(@RequestParam(value = "message",
            defaultValue = "你好，请问你的知识库文档主要是关于什么内容的?") String message, HttpServletResponse response) {
        response.setCharacterEncoding("UTF-8");
        Query transform = queryTransformer.transform(new Query(message));
        String text = transform.text();
        System.out.println("query:" + text);

        return cloudRagService.retrieve(message).map(x -> x.getResult().getOutput().getText());
    }

}
